package com.tiantianfund.udf;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.*;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.StringObjectInspector;
import org.apache.hadoop.io.Text;
import org.json.JSONArray;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.ArrayList;
import java.util.List;

public class GetJsonArray extends GenericUDTF {
    private static final Logger LOG = LoggerFactory.getLogger(GenericUDTF.class.getName());

    private transient ObjectInspector inputOIs; // input ObjectInspectors
    private transient final Object[] forwardObj = new Object[1];
    private final Text result = new Text();

    @Override
    public StructObjectInspector initialize(StructObjectInspector argOIs) throws UDFArgumentException {
        /***
         * 这里直接抄源码里已经写好的UDTF函数即可
         * 比如：GenericUDTFExplode，GenericUDTFJSONTuple，GenericUDTFGetSplits
         *
         */
        List<? extends StructField> inputFields = argOIs.getAllStructFieldRefs();
        if (inputFields.size() != 1) {
            throw new UDFArgumentException("get_json_array() takes only one argument");
        }

        inputOIs = inputFields.get(0).getFieldObjectInspector();

        ArrayList<String> fieldNames = new ArrayList<String>(1);
        ArrayList<ObjectInspector> fieldOIs = new ArrayList<ObjectInspector>(1);
        //添加col_name1
        fieldNames.add("col_name1");
        fieldOIs.add(PrimitiveObjectInspectorFactory.writableStringObjectInspector);
        return ObjectInspectorFactory.getStandardStructObjectInspector(fieldNames, fieldOIs);
    }

    @Override
    public void process(Object[] args) throws HiveException {
        Object odsJson = args[0];
        // 这里必须判断null，否则直接报错，程序中传入null应该属于正常，所以返回空。
        if (odsJson == null) {
            return;
        }

        // 空json数组，返回空字符串
        String newJson = ((StringObjectInspector) inputOIs).getPrimitiveJavaObject(odsJson);
        JSONArray jsonArray = new JSONArray(newJson);
        if (jsonArray.length() == 0) {
            result.set("");
            forwardObj[0] = result;
            forward(forwardObj);
        }

        for (int i = 0; i < jsonArray.length(); i++) {
            result.set(jsonArray.get(i).toString()); // 转成text再返回，很重要
            forwardObj[0] = result;
            forward(forwardObj);
        }
    }

    @Override
    public void close() throws HiveException {

    }

}
